from torch.nn import Conv2d, Sequential, ModuleList, ReLU
from detection.models.yolo.yolox.yolo_pafpn import YOLOPAFPN
from detection.models.yolo.yolox.network_blocks import BaseConv, DWConv
import torch.nn as nn
import torch
import torch.nn.functional as F
from detection.anchor_utils import anchor_utils


def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0):
    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
    """
    return Sequential(
        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
               groups=in_channels, stride=stride, padding=padding),
        ReLU(True),
        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
    )


def HeadConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0):
    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
    """
    return Sequential(
        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
               stride=stride, padding=padding),
        ReLU(True),
        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
    )


class YOLOSSD(nn.Module):
    def __init__(self,
                 num_classes,
                 backbone,
                 bbox_headers,
                 class_headers,
                 is_test=False,
                 prior_boxes=None,
                 device=None):
        super(YOLOSSD, self).__init__()

        self.backbone = backbone
        self.bbox_headers = bbox_headers
        self.class_headers = class_headers
        self.num_classes = num_classes
        self.is_test = is_test
        self.device = device

        self.center_variance = prior_boxes.center_variance
        self.size_variance = prior_boxes.size_variance
        self.priors = prior_boxes.priors.to(self.device)
        self.freeze_header = prior_boxes.freeze_header

    def forward(self, x):
        confidences = []
        locations = []
        features = self.backbone(x)
        for i in range(len(features)):
            conf, loc = self.compute_header(i, features[i])
            confidences.append(conf)
            locations.append(loc)
        confidences = torch.cat(confidences, 1)
        locations = torch.cat(locations, 1)
        if self.is_test:
            confidences = F.softmax(confidences, dim=2)
            if self.freeze_header:
                locations = anchor_utils.decode(locations, self.priors, [self.center_variance, self.size_variance])
        return locations, confidences

    def compute_header(self, i, x):
        """
        input:320
        x=torch.Size([1, 64, 40, 40]),location:torch.Size([1, 12, 40, 40])
        x=torch.Size([1, 128, 20, 20]),location:torch.Size([1, 8, 20, 20])
        x=torch.Size([1, 256, 10, 10]),location:torch.Size([1, 8, 10, 10])
        x=torch.Size([1, 256, 5, 5]),location:torch.Size([1, 12, 5, 5])
        :param i:
        :param x:
        :return:
        """
        conf = self.class_headers[i](x)
        conf = conf.permute(0, 2, 3, 1).contiguous()
        conf = conf.view(conf.size(0), -1, self.num_classes)

        loc = self.bbox_headers[i](x)
        # print("x={},location:{}".format(x.shape, loc.shape))
        loc = loc.permute(0, 2, 3, 1).contiguous()
        loc = loc.view(loc.size(0), -1, 4)
        return conf, loc


def create_model_yolox_ssd(prior_boxes,
                           num_classes,
                           is_test=False,
                           width_mult=1.0,
                           pretrained=False,
                           device="cuda:0",
                           net_type="yolox-nano",
                           **kwargs):
    """
    https://github.com/Megvii-BaseDetection/YOLOX
    yolox-nano, yolox-tiny, yolox-s, yolox-m(64), yolox-l, yolox-x
    :param prior_boxes:
    :param num_classes:
    :param is_test:
    :param width_mult:
    :return:
    """
    # feature size=[40,20,10]
    param_dict = {
        "yolox-nano": (0.33, 0.25),
        "yolox-tiny": (0.33, 0.375),
        "yolox-s": (0.33, 0.50),
        "yolox-m": (0.67, 0.75),
        "yolox-l": (1.0, 1.0),
        "yolox-x": (1.33, 1.25),
    }
    if not net_type in param_dict:
        raise Exception(" YOLOX only support net_type={},not support {}".format(list(param_dict.keys()), net_type))
    depth, width = param_dict[net_type]
    backbone = YOLOPAFPN(depth, width, width_mult=width_mult)
    channels = backbone.get_head_channels()
    boxes_expand = [len(boxes) * (len(prior_boxes.aspect_ratios)) for boxes in prior_boxes.min_sizes]
    bbox_headers = ModuleList([SeperableConv2d(in_channels=channels[0],
                                               out_channels=boxes_expand[0] * 4,
                                               kernel_size=3,
                                               padding=1),
                               SeperableConv2d(in_channels=channels[1],
                                               out_channels=boxes_expand[1] * 4,
                                               kernel_size=3,
                                               padding=1),
                               Conv2d(in_channels=channels[2],
                                      out_channels=boxes_expand[2] * 4,
                                      kernel_size=3,
                                      padding=1)
                               ])

    class_headers = ModuleList([SeperableConv2d(in_channels=channels[0],
                                                out_channels=boxes_expand[0] * num_classes,
                                                kernel_size=3,
                                                padding=1),
                                SeperableConv2d(in_channels=channels[1],
                                                out_channels=boxes_expand[1] * num_classes,
                                                kernel_size=3,
                                                padding=1),
                                Conv2d(in_channels=channels[2],
                                       out_channels=boxes_expand[2] * num_classes,
                                       kernel_size=3,
                                       padding=1)
                                ])

    return YOLOSSD(num_classes,
                   backbone,
                   bbox_headers,
                   class_headers,
                   is_test,
                   prior_boxes,
                   device)


def create_model_yolox_v2_ssd(prior_boxes,
                              num_classes,
                              is_test=False,
                              width_mult=1.0,
                              device="cuda:0",
                              net_type="yolox-nano",
                              **kwargs):
    """
    https://github.com/Megvii-BaseDetection/YOLOX
    yolox-nano, yolox-tiny, yolox-s, yolox-m(64), yolox-l, yolox-x
    :param net_type: {"yolox-nano","yolox-tiny","yolox-s","yolox-m","yolox-l","yolox-x",}
    :param depthwise: True or False
    :param act: {"silu","relu","lrelu"}
    """
    param_dict = {
        "yolox-nano": (0.33, 0.25),
        "yolox-tiny": (0.33, 0.375),
        "yolox-s": (0.33, 0.50),
        "yolox-m": (0.67, 0.75),
        "yolox-l": (1.0, 1.0),
        "yolox-x": (1.33, 1.25),
    }
    if not net_type in param_dict:
        raise Exception(" YOLOX only support net_type={},not support {}".format(list(param_dict.keys()), net_type))
    depthwise = kwargs["depthwise"] if "depthwise" in kwargs else False
    # act={"silu","relu","lrelu"}
    act = kwargs["act"] if "act" in kwargs else "silu"
    depth, width = param_dict[net_type]
    backbone = YOLOPAFPN(depth, width, width_mult=width_mult, depthwise=depthwise, act=act)
    channels = backbone.get_head_channels()
    boxes_expand = [len(boxes) * (len(prior_boxes.aspect_ratios)) for boxes in prior_boxes.min_sizes]
    # Conv = DWConv if depthwise else BaseConv
    # Conv = SeperableConv2d if depthwise else HeadConv2d
    bbox_headers = ModuleList([SeperableConv2d(in_channels=channels[0],
                                               out_channels=boxes_expand[0] * 4,
                                               kernel_size=3,
                                               padding=1),
                               SeperableConv2d(in_channels=channels[1],
                                               out_channels=boxes_expand[1] * 4,
                                               kernel_size=3,
                                               padding=1),
                               Conv2d(in_channels=channels[2],
                                      out_channels=boxes_expand[2] * 4,
                                      kernel_size=3,
                                      padding=1)
                               ])

    class_headers = ModuleList([SeperableConv2d(in_channels=channels[0],
                                                out_channels=boxes_expand[0] * num_classes,
                                                kernel_size=3,
                                                padding=1),
                                SeperableConv2d(in_channels=channels[1],
                                                out_channels=boxes_expand[1] * num_classes,
                                                kernel_size=3,
                                                padding=1),
                                Conv2d(in_channels=channels[2],
                                       out_channels=boxes_expand[2] * num_classes,
                                       kernel_size=3,
                                       padding=1)
                                ])

    return YOLOSSD(num_classes,
                   backbone,
                   bbox_headers,
                   class_headers,
                   is_test,
                   prior_boxes,
                   device)